from fastapi import FastAPI, Query, HTTPException
from pydantic import BaseModel
from handlers.mongo_handler import MongoDBHandler
from handlers.qdrant_handler import QdrantHandler
from config import DB_CONFIG, Qdrant_CONFIG

mongo_handler = MongoDBHandler(DB_CONFIG['url'])
qdrant_handler = QdrantHandler(Qdrant_CONFIG['url'], Qdrant_CONFIG['collection_name'], Qdrant_CONFIG['embedding_model'])

app = FastAPI()

class SearchRequest(BaseModel):
    db_name: str
    collection_name: str
    exclude_fields: list = None

class InsertRequest(BaseModel):
    documents: list

class SearchQdrantRequest(BaseModel):
    keyword: str
    top_k: int = 5

@app.get("/documents")
async def get_documents(
    db_name: str = Query(DB_CONFIG["db_name"]),
    collection_name: str = Query(DB_CONFIG["collection_name"]),
):
    results = mongo_handler.search_documents(db_name, collection_name)
    return {"results": results}

@app.post("/insert_documents")
async def insert_documents(request: InsertRequest):
    try:
        qdrant_handler.insert_documents(request.documents)
        return {"message": "Documents inserted successfully"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/search")
async def search_qdrant(
    query: str = Query(..., description="查询关键词"),
    top_k: int = Query(5, description="返回前K条结果")
):
    try:
        results = qdrant_handler.search(query, top_k)
        return {"results": results}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
